package wmi

import (
	"fmt"
	"github.com/yusufpapurcu/wmi"
)

// WbemTest
// https://learn.microsoft.com/zh-cn/windows/win32/cimwin32prov/computer-system-hardware-classes
type WbemTest struct {
}

func NewWbemTest() *WbemTest {
	return &WbemTest{}
}

// GetProcess 进程
func (tis *WbemTest) GetProcess(ProcessId ...int) ([]Win32_Process, error) {
	wql := fmt.Sprintf("SELECT * FROM Win32_Process")
	if len(ProcessId) > 0 {
		wql = fmt.Sprintf("SELECT * FROM Win32_Process WHERE ProcessId=%v", ProcessId[0])
	}

	var values []Win32_Process
	if err := wmi.Query(wql, &values); err != nil {
		return nil, err
	}

	return values, nil
}

// GetProcessBy 进程
func (tis *WbemTest) GetProcessBy(name string, pid int) ([]Win32_Process, error) {
	wql := fmt.Sprintf("SELECT * FROM Win32_Process ")
	if len(name) > 0 || pid > 0 {
		hasFirst := false
		wql += " WHERE "
		if pid > 0 {
			wql += fmt.Sprintf(" ProcessId=%v ", pid)
			hasFirst = true
		}
		if len(name) > 0 {
			if hasFirst {
				wql += " AND "
			}
			wql += fmt.Sprintf(" (Caption=\"%v\" OR Caption=\"%v.exe\") ", name, name)
		}
	}

	var values []Win32_Process
	if err := wmi.Query(wql, &values); err != nil {
		return nil, err
	}

	return values, nil
}

// GetPhysicalMemory 物理内存
func (tis *WbemTest) GetPhysicalMemory() ([]Win32_PhysicalMemory, error) {
	var values []Win32_PhysicalMemory
	if err := wmi.Query("SELECT * FROM Win32_PhysicalMemory", &values); err != nil {
		return nil, err
	}

	return values, nil
}

// GetProcessor 处理器
func (tis *WbemTest) GetProcessor() ([]Win32_Processor, error) {
	var values []Win32_Processor
	if err := wmi.Query("SELECT * FROM Win32_Processor", &values); err != nil {
		return nil, err
	}

	return values, nil
}

// GetDiskDrive 硬盘驱动器
func (tis *WbemTest) GetDiskDrive() ([]Win32_DiskDrive, error) {
	var values []Win32_DiskDrive
	if err := wmi.Query("SELECT * FROM Win32_DiskDrive", &values); err != nil {
		return nil, err
	}

	return values, nil
}

// GetLogicalDisk 逻辑磁盘
func (tis *WbemTest) GetLogicalDisk() ([]Win32_LogicalDisk, error) {
	var values []Win32_LogicalDisk
	if err := wmi.Query("SELECT * FROM Win32_LogicalDisk", &values); err != nil {
		return nil, err
	}

	return values, nil
}

// GetDiskPartition 物理磁盘分区
func (tis *WbemTest) GetDiskPartition() ([]Win32_DiskPartition, error) {
	var values []Win32_DiskPartition
	if err := wmi.Query("SELECT * FROM Win32_DiskPartition", &values); err != nil {
		return nil, err
	}

	return values, nil
}

// GetLogicalDiskWithDiskDrive 逻辑磁盘
func (tis *WbemTest) GetLogicalDiskWithDiskDrive() ([]*LogicalDisk, error) {
	logicalDisk, err := tis.GetLogicalDisk()
	if err != nil {
		return nil, err
	}

	var result = map[string]*LogicalDisk{}
	for _, item := range logicalDisk {
		result[item.Name] = &LogicalDisk{
			Win32_LogicalDisk: item,
			DiskDrive:         nil,
		}
	}

	diskDrive, err := tis.GetDiskDrive()
	if err != nil {
		return nil, err
	}

	for _, disk := range diskDrive {
		wql := fmt.Sprintf("ASSOCIATORS OF {Win32_DiskDrive.DeviceID='%v'} WHERE AssocClass = Win32_DiskDriveToDiskPartition", disk.DeviceID)

		var partitions []Win32_DiskPartition
		if err = wmi.Query(wql, &partitions); err != nil {
			return nil, err
		}

		for _, partition := range partitions {
			wql = fmt.Sprintf("ASSOCIATORS OF {Win32_DiskPartition.DeviceID='%v'} WHERE AssocClass = Win32_LogicalDiskToPartition", partition.DeviceID)
			var tmpLogicalDisk []Win32_LogicalDisk
			if err = wmi.Query(wql, &tmpLogicalDisk); err != nil {
				return nil, err
			}

			for _, item := range tmpLogicalDisk {
				if object, ok := result[item.Name]; ok {
					tmpDist := disk
					object.DiskDrive = &tmpDist
				}
			}
		}

	}

	var values []*LogicalDisk
	for _, item := range result {
		object := item
		values = append(values, object)
	}

	return values, err
}

// GetOperatingSystem 操作系统信息
func (tis *WbemTest) GetOperatingSystem() (*Win32_OperatingSystem, error) {
	var values []Win32_OperatingSystem
	if err := wmi.Query("SELECT * FROM Win32_OperatingSystem", &values); err != nil {
		return nil, err
	}

	if len(values) == 0 {
		return nil, fmt.Errorf("not found value")
	}

	return &values[0], nil
}
